import requests
import pandas as pd
import datetime as dt
import re
from numpy import NaN


def all_genes_one_by_one():
    for i in range(1, 10):
        url = "https://push2.eastmoney.com/api/qt/clist/get?" \
              "np=1&fltt=2" \
              "&invt=2" \
              "&fields=f1,f2,f3,f4,f12,f13,f14,f128" \
              "&pn=%d&pz=30&fid=f3" \
              "&po=1&fs=m:90+t:3" \
              "&forcect=1&_=07512" % i
        resp = requests.get(url)
        if resp.ok:
            resp_json = resp.json()
            resp_data = resp_json['data']['diff']
            for data in resp_data:
                yield {'code': data['f12'],
                       'name': data['f14'],
                       'market': data['f13'],
                       'latest_price': data['f2'],
                       'raise_price': data['f4'],
                       'raise_range': data['f3'],
                       'stock_code': data['f140'],
                       'stock_name': data['f128'],
                       'stock_market': data['f141']
                       }


def all_genes():
    datas = [data for data in all_genes_one_by_one()]
    df = pd.DataFrame(datas)
    return df


def _part_gene_klines(gid, market=90, end=dt.date.today()):
    if isinstance(end, dt.date):
        end = end.strftime('%Y%m%d')
    url = "https://push2his.eastmoney.com/api/qt/stock/kline/get?" \
          "secid={market}.{g_id}&klt=101&fqt=0&lmt=80&end={end}&iscca=1&fields1=f1%2Cf2%2Cf3%2Cf4%2Cf5" \
          "&fields2=f51%2Cf52%2Cf53%2Cf54%2Cf55%2Cf56%2Cf57" \
          "&ut=f057cbcbce2a86e2866ab8877db1d059" \
          "&forcect=1&fsdata1581835516=fsdata1581835516".replace('{market}', str(market)).replace('{g_id}',
                                                                                                  gid).replace('{end}',
                                                                                                               end)
    resp = requests.get(url)
    if not resp.ok:
        return [], 0
    else:
        print(resp.json())
        data = resp.json()['data']
        if data:
            klines = data['klines']
            results = [{
                    "code": gid,
                    "date": _kl[0],
                    "open": _kl[1],
                    "close": _kl[2],
                    "high": _kl[3],
                    "low": _kl[4],
                    'total_volume': _kl[5],
                    'total_amount': _kl[6]
                } for _kl in [
                kl.split(',') for kl in klines
            ]]
            # 涨幅价格
            # 总数也返回
            return results, data['dktotal']
        else:
            return [], 0


def load_gene_klines(gid, market=90, end=dt.date.today()):
    if isinstance(end, dt.date):
        end = end.strftime('%Y%m%d')
    url = "https://push2his.eastmoney.com/api/qt/stock/kline/get?" \
          "secid={market}.{g_id}&klt=101&fqt=0&lmt=80&end={end}&iscca=1&fields1=f1%2Cf2%2Cf3%2Cf4%2Cf5" \
          "&fields2=f51%2Cf52%2Cf53%2Cf54%2Cf55%2Cf56%2Cf57" \
          "&ut=f057cbcbce2a86e2866ab8877db1d059" \
          "&forcect=1&fsdata1581835516=fsdata1581835516".replace('{market}', str(market)).replace('{g_id}',
                                                                                                  gid).replace('{end}',
                                                                                                               end)
    resp = requests.get(url)
    if not resp.ok:
        raise Exception('load error with code %d' % resp.status_code)
    else:
        data = resp.json()['data']
        klines = data['klines']
        dfs = ([], [], [], [], [],)
        for kl in klines:
            _kl = kl.split(',')
            for i in range(0, len(_kl)):
                if i < len(dfs):
                    if i == 0:
                        dfs[i].append(_kl[i])
                    else:
                        dfs[i].append(float(_kl[i]))
        dates, Opens, Closes, Highs, Lows = dfs
        df = pd.DataFrame(data={
            'dates': dates,
            'Open': Opens,
            'Close': Closes,
            'High': Highs,
            'Low': Lows
        })
        df['code'] = data['code']
        df['name'] = data['name']
        # 涨幅价格
        # 总数也返回
        return df, data['dktotal']


def load_gene_all_klines(gid, market=90):
    datas = []
    s = 0
    end = dt.date.today()
    oneday = dt.timedelta(days=1)
    while True:
        data, total = load_gene_klines(gid, market, end)
        datas.insert(0, data)
        s += len(data)
        if s >= total or len(data) == 0:
            break
        else:
            first = data['dates'].head(1).values[0]
            first_day = dt.datetime.strptime(first, '%Y-%m-%d')
            end = first_day - oneday
    return pd.concat(datas) \
        .assign(Raise=lambda x: x['Close'].rolling(window=2)
                .apply(lambda y: round(100 * y[1] / y[0] - 100, 2), raw=True))


def load_gene_all_klines_one_by_one(gid, market=90, start=None):
    end = dt.date.today()
    s = 0
    oneday = dt.timedelta(days=1)
    while True:
        data, total = _part_gene_klines(gid, market, end)
        s += len(data)
        if s >= total or len(data) == 0:
            break
        else:
            for i in range(len(data) - 1, -1, -1):
                yield data[i]
            first = data[0]['date']
            first_day = dt.datetime.strptime(first, '%Y-%m-%d')
            if start is not None and first_day <= start:
                break
            end = first_day - oneday


def load_gene_stocks_one_by_one(gid):
    size = 20
    page = 1
    while True:
        url = "https://push2.eastmoney.com/api/qt/clist/get?" \
              "ut=f057cbcbce2a86e2866ab8877db1d059&forcect=1" \
              "&fs=b:%s&pn=%d&pz=%d&po=1&fid=f3&invt=2" % (gid, page, size)
        resp = requests.get(url)
        if not resp.ok:
            break
        else:
            print(url)
            data = resp.json()['data']
            diff = data['diff']
            for k in diff.keys():
                stock = diff[k]
                yield {
                    'code': stock['f12'],
                    'name': stock['f14'],
                    'market': str(stock['f1']),
                    'latest': stock['f2'],
                    'raise': stock['f3']
                }
            total = data['total']
            if size * page >= total:
                break
            else:
                page = page + 1


def load_gene_stocks(gid):
    return None


def get_stock_sec(code):
    resp = requests.get('https://push2.eastmoney.com/api/qt/stock/details/get'
                        '?secid=%s&ut=f057cbcbce2a86e2866ab8877db1d059&forcect=1'
                        '&fields1=f1,f2,f3,f4,f5&fields2=f51,f52,f53,f54,f55'
                        '&pos=-14&iscca=1&invt=2' % code)
    if resp.ok:
        data = resp.json()
        return data
    else:
        return None


if __name__ == '__main__':
    for data in load_gene_all_klines_one_by_one('BK0711'):
        print(data)
